# -*- coding: utf-8 -*-
import json
import os
from tqdm import tqdm
import copy


def labelme(json_path, save_path, point_count, add_visible):
    d = json.load(open(json_path))
    imgh = d["imageHeight"]
    imgw = d["imageWidth"]

    # 存储一张图片中所有的目标
    labels_dict = {}
    # 一个目标作为一个组
    obj_dict = {"label": "", "rectangle": [], "keypoints": [[] for i in range(point_count)]}
    for target in d["shapes"]:
        # 目标的编号作为键
        group_id = str(target["group_id"])
        shape_type = target["shape_type"]

        # 没有这个目标(组)就创建
        if group_id not in labels_dict:
            # 深拷贝，防止修改obj_dict模板  使用.copy()，内部嵌套的对象还是浅拷贝，比如keypoints
            labels_dict[group_id] = copy.deepcopy(obj_dict)

        # 是坐标框，添加目标标签(比如人)和坐标框(x1, y1, x2, y2)
        if shape_type == "rectangle":
            label = target["label"]
            points = target["points"]
            x1, y1, x2, y2 = points[0][0], points[0][1], points[1][0], points[1][1]

            labels_dict[group_id]["label"] = label
            labels_dict[group_id]["rectangle"] = [x1, y1, x2, y2]

        # 是关键点，按关键点名称(0,1,2...)添加坐标
        elif shape_type == "point":
            label = target["label"]
            points = target["points"]
            if add_visible:
                # 描述信息内写关键点是否可见(0,1,2),0：不可见，1：被遮挡，2：可见
                visible = int(target["description"])
                labels_dict[group_id]["keypoints"][int(label)] = [points[0][0], points[0][1], visible]
            else:
                labels_dict[group_id]["keypoints"][int(label)] = [points[0][0], points[0][1]]

    # save txt path
    txt_file = open(save_path, 'w')
    print(labels_dict)

    for key, obj in labels_dict.items():
        label = obj['label']
        label_index = clas_name.index(label)
        box = obj['rectangle']
        x1, y1, x2, y2 = box
        cx = ((x1 + x2) / 2) / imgw
        cy = ((y1 + y2) / 2) / imgh
        w = (x2 - x1) / imgw
        h = (y2 - y1) / imgh

        keypoints = obj['keypoints']

        point_list = []
        for point in keypoints:
            # 没有标注这个关键点报错
            if not point:
                print('\033[31m' + f"Path :{json_path} is Error! Please check the annotation file" + '\033[0m')
                exit()

            if add_visible:
                x, y, v = point
            else:
                x, y = point
            point_list.append(x / imgw)
            point_list.append(y / imgh)
            if add_visible:
                point_list.append(v)

        point_list = list(map(lambda x: str(x), point_list))
        point_list_str = ' '.join(point_list)
        txt_file.write(f"{label_index} {cx} {cy} {w} {h} {point_list_str}" + "\n")

    txt_file.close()


if __name__ == "__main__":
    # 标注时，[需要标注坐标框，标注关键点](每个目标一个Group ID)。[每个关键点的Label description信息内填是否可见](0,1,2)

    jsons_path = r".\data\annotations"
    save_dir = r".\data\labels"
    # images_path = r"\data\images"

    clas_name = ["person"]
    # 单个目标关键点个数
    point_count = 17
    # 是否添加可见性参数，添加一个关键3个值，不添加，一个关键2个值
    add_visible = True

    files = os.listdir(jsons_path)
    json_files = [file for file in files if file.endswith('.json')]
    for json_name in tqdm(json_files):
        # print(json_name)
        if json_name[0] == ".":
            continue
        json_path = os.path.join(jsons_path, json_name)
        save_path = os.path.join(save_dir, json_name.split(".")[0] + ".txt")
        # print(json_path)
        # print(save_path)

        labelme(json_path, save_path, point_count, add_visible)
